{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Add message history (memory)\n",
    "\n",
    "The RunnableWithMessageHistory lets us add message history to certain types of chains. It wraps another Runnable and manages the chat message history for it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "from langchain_openai.chat_models import ChatOpenAI\n",
    "\n",
    "model = ChatOpenAI()\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"You're an assistant who's good at {ability}. Respond in 20 words or fewer\",\n",
    "        ),\n",
    "        MessagesPlaceholder(variable_name=\"history\"),\n",
    "        (\"human\", \"{input}\"),\n",
    "    ]\n",
    ")\n",
    "runnable = prompt | model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## In-memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.chat_message_histories import ChatMessageHistory\n",
    "from langchain_core.chat_history import BaseChatMessageHistory\n",
    "from langchain_core.runnables.history import RunnableWithMessageHistory\n",
    "\n",
    "store = {}\n",
    "\n",
    "\n",
    "def get_session_history(session_id: str) -> BaseChatMessageHistory:\n",
    "    if session_id not in store:\n",
    "        store[session_id] = ChatMessageHistory()\n",
    "    return store[session_id]\n",
    "\n",
    "\n",
    "with_message_history = RunnableWithMessageHistory(\n",
    "    runnable,\n",
    "    get_session_history,\n",
    "    input_messages_key=\"input\",\n",
    "    history_messages_key=\"history\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='Cosine is a trigonometric function that represents the ratio of the adjacent side to the hypotenuse in a right triangle.', response_metadata={'token_usage': {'completion_tokens': 26, 'prompt_tokens': 33, 'total_tokens': 59}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-d32f7e4f-33f6-4777-ad32-2808c6311c52-0', usage_metadata={'input_tokens': 33, 'output_tokens': 26, 'total_tokens': 59})"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with_message_history.invoke(\n",
    "    {\"ability\": \"math\", \"input\": \"What does cosine mean?\"},\n",
    "    config={\"configurable\": {\"session_id\": \"abc123\"}},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='Cosine tells us the horizontal distance compared to the hypotenuse in a right triangle.', response_metadata={'token_usage': {'completion_tokens': 18, 'prompt_tokens': 69, 'total_tokens': 87}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-e2c8007d-02fc-4773-8f1f-b190360dcab1-0', usage_metadata={'input_tokens': 69, 'output_tokens': 18, 'total_tokens': 87})"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Remembers\n",
    "with_message_history.invoke(\n",
    "    {\"ability\": \"math\", \"input\": \"What?\"},\n",
    "    config={\"configurable\": {\"session_id\": \"abc123\"}},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='I can assist you with any math questions you have. Just let me know what you need help with!', response_metadata={'token_usage': {'completion_tokens': 21, 'prompt_tokens': 30, 'total_tokens': 51}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-bb0a76a7-e44f-4f3c-a263-d0f05adf112d-0', usage_metadata={'input_tokens': 30, 'output_tokens': 21, 'total_tokens': 51})"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# New session_id --> does not remember.\n",
    "with_message_history.invoke(\n",
    "    {\"ability\": \"math\", \"input\": \"What?\"},\n",
    "    config={\"configurable\": {\"session_id\": \"def234\"}},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Persistent storage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.chat_message_histories import RedisChatMessageHistory\n",
    "\n",
    "# pip install --upgrade --quiet redis\n",
    "def get_message_history(session_id: str) -> RedisChatMessageHistory:\n",
    "    return RedisChatMessageHistory(session_id, url=\"redis://localhost:6379/0\")\n",
    "\n",
    "\n",
    "with_message_history = RunnableWithMessageHistory(\n",
    "    runnable,\n",
    "    get_message_history,\n",
    "    input_messages_key=\"input\",\n",
    "    history_messages_key=\"history\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with_message_history.invoke(\n",
    "    {\"ability\": \"math\", \"input\": \"What does cosine mean?\"},\n",
    "    config={\"configurable\": {\"session_id\": \"foobar\"}},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with_message_history.invoke(\n",
    "    {\"ability\": \"math\", \"input\": \"What's its inverse\"},\n",
    "    config={\"configurable\": {\"session_id\": \"foobar\"}},\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain0_1",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
